#!/usr/bin/python
# GM_SimpleIntegrate.py

#import math
import numpy
from time import asctime
from os import linesep

# Must have G, c in appropriate units so rE is in pc! Note M is in Msun
# See notes pg. 13 date 2/22/11
Goverc2 = 4.786e-14

VariableNames = ['appearance', 'masstuple', 'vtuple', 'd', 'deltad', 'ba', 'deltaba', 'deltala', 'MvVals', 'delMvVals', 'MvMass', 'delMvMass', 'dlocs', 'blocs', 'SourceMag', 'delMvList', 'darkstars']


def GetEinsteinModifier(Mv):
    return 1.

def CalculateMCparams(simpars, phypars, nMv, nM):
    for i in simpars:
        cmd = "%s = simpars['%s']" % (i,i)
        exec cmd
    for i in phypars:
        cmd = "%s = phypars['%s']" % (i,i)
        exec cmd
    # The minimimum number of draws from the random generator is the number of integration elements: 0.5*lres*bres*dres*(dres + 1) (the number of spatial elements including a convolution along one of the axes; if pencilBeam == True, lres and bres are going to get set to one later, and that is reflected here) times nMv (the resolution in magnitude) times nM (the resolution in mass).
    if pencilBeam == False:
        NtMin = numpy.ceil(0.5*lres*bres*dres*(dres + 1)*nMv*nM)
    elif pencilBeam == True:
        NtMin = numpy.ceil(0.5*dres*(dres + 1)*nMv*nM)
    # The "padded requested number" multiplies by a buffer factor intended to account for the fact that some draws on the random number generator will get thrown out.  This is additionally padded with a five-sigma "root N" Poisson width.  If your guess for bufferfactor was reasonable, you should (pretty much) always get enough timescales.  If you don't get enough, increase bufferfactor and try again.
    PaddedRequest = NtRequested + 5*NtRequested**0.5
    PaddedRequest = numpy.ceil(PaddedRequest*bufferfactor)
    if PaddedRequest <= NtMin:
        NtMult = 1
        NtGiven = int(NtMin)
    elif PaddedRequest > NtMin:
        NtMult = int(numpy.ceil(PaddedRequest/float(NtMin)))
        NtGiven = int(NtMult*NtMin)
#    print "lres, bres, dres, len(MvVals), len(MvMass))"
#    print lres, bres, dres, len(MvVals), len(MvMass)
#    print "nMv, nM"
#    print nMv, nM
#    print "NtMin, PaddedRequest, NtMult, NtGiven"
#    print NtMin, PaddedRequest, NtMult, NtGiven
    for i in simpars:
        cmd = "simpars['%s'] = %s" % (i,i)
        exec cmd
    return simpars

def simple(simpars, phypars, appearance, masstuple, vtupleSource, vtupleLens):
    '''Acquire lensing rates, qualified only by source magnitude.'''
    # First, fix the MC paramaters.  I would do this at the standard spot (GM_alterpars), but for some reason something was wrong there, and I am feeling too lazy to sort it out when this solution works.
    simpars = CalculateMCparams(simpars, phypars, len(appearance), len(masstuple[0]))
    # Intake 'n' stuff
    for i in simpars:
        cmd = "%s = simpars['%s']" % (i,i)
        exec cmd
    for i in phypars:
        cmd = "%s = phypars['%s']" % (i,i)
        exec cmd
    if len(appearance) == 1:
        print "Process", codeindex, "beginning."
        print asctime()
    (darkstars, rho0mlist) = masstuple
    (vrelappl_S, sigmaappl2_S, sigmaappb2_S) = vtupleSource
    (vrelappl_L, sigmaappl2_L, sigmaappb2_L) = vtupleLens
    # Distance elements
    DL = numpy.array(d)
    dDL = numpy.array(deltad)
    DLdDL = DL*dDL
    DS = numpy.array(d)
    dDS = numpy.array(deltad)
    DS2dDS = DS*DS*dDS
    # Angular element, incl. symmetry restoration via quadrupling
    cosb = numpy.cos(ba)
    cosbdeltab = cosb*deltaba
    dOmega = numpy.multiply(deltala[:,numpy.newaxis], cosbdeltab[numpy.newaxis,:])*4 # Symmetry!
    del cosb, cosbdeltab
    # These histograms compile things already "naturally" separated and hold a running tally at any point in the program.  DSHist, DLHist, and bHist are collected before final summation for each 
    MvHist = numpy.zeros(len(delMvVals), dtype=float) # Divide by dMv here
    MassHist = numpy.zeros(len(MvMass), dtype=float) # Divide by dM here
    DSHist = numpy.zeros(len(dlocs), dtype=float) # Divide by dDS here
    DLHist = numpy.zeros(len(dlocs), dtype=float) # Divide by dDL here
    bHist = numpy.zeros(len(blocs), dtype=float) # Divide by dba here
    # These "master histograms" contain snapshots of the above histograms at each of the limiting source magnitude cutoffs.  Each falls along a natural dimension of the problem, so to speak, along a quantity that needed to be integrated over to start with.  A sixth dimension, l, could be added here, but has not been called for in any way, so I leave it out.
    MasterMvHist = []
    MasterMassHist = []
    MasterDSHist = []
    MasterDLHist = []
    MasterbHist = []
    # These two arrays store a flattened-out copy of lensing rate contribution (which is used as a probabilistic weight later) and draws from the timescale distribution at the corresponding locations.
    timescaleArray = numpy.zeros(NtGiven)
    timescaleWeight = numpy.zeros(NtGiven)
    # These histograms need additional complex masks to delimit them
#    vtappHist =  numpy.ones(110, dtype=float) #### Not supported yet
    # lensrates holds the overall result
    lensrate = 0.0
    lensrates = numpy.zeros(len(SourceMag))
    print "len(SourceMag)", len(SourceMag)
    SourceMagIndex = 0
    CurrentLimitingMag = SourceMag[SourceMagIndex]
    timescaleMasterIndex = 0
    # i indexes source apparent magnitude
    for i in range(0, len(appearance)):
        nS = numpy.array(appearance[i])
        Mv = float(MvVals[i]) # Intrinsic magnitudes are from MvList, but apparent magnitudes are from MvVals
        dMv = float(delMvList[i])
        EinsteinModifier = GetEinsteinModifier(Mv) # A not-yet-implemented sensitivity-as-a-funtion-of-magnitude modifier
        # Cut for source unlensed magnitude
        # This section (in the if-statement) is kind of the end of the previous loop, rather than the beginning of this loop, just... oddly placed.  For reasons that totally made sense at the time, somehow.
        if Mv > (CurrentLimitingMag + 0.00001):
            lensrates[SourceMagIndex] = lensrate
            MasterMvHist.append(numpy.array(MvHist))
            MasterMassHist.append(numpy.array(MassHist))
            MasterDSHist.append(numpy.array(DSHist))
            MasterDLHist.append(numpy.array(DLHist))
###
            # If you've reached the end, stop; otherwise, move to next stage
            if SourceMagIndex == (len(SourceMag) - 1):
                break
            else:
                SourceMagIndex = SourceMagIndex + 1
                CurrentLimitingMag = SourceMag[SourceMagIndex]
###
        icontrib = 0.0
        # j indexes lens mass
        for j in range(0, len(darkstars)):
            nL = darkstars[j]
            M = MvMass[j,1]
            ijcontrib = 0.0
            # k indexes over distance from earth, or more accurately, over the value of DS - DL
            # Source quantities get sliced [Ki:,:,:]
            # Lens quantities get sliced [:Kf,:,:]
            # Hence when they're combined, DS > DL always
            for k in range(0, len(dlocs)):
                Ki = k
                Kf = len(dlocs) - k
                # 1d, [d]
                # two_rE:  2x Einstein radius, times a modifier for photometric sensitivity, if available; parsecs
                # DLdDL_DS2dDS:  Volume elements & additional distance terms to correct
                # AllDistances:  Combine two_rE and DLdDL_DS2dDS
                # EinsteinAngle:  Einstein angle, radians, for use in velocity Monte Carlo stuff
                two_rE = 2*(4*Goverc2*M*DL[:Kf]*(DS[Ki:] - DL[:Kf])/DS[Ki:])**0.5*EinsteinModifier
                DLdDL_DS2dDS = DLdDL[:Kf]*DS2dDS[Ki:]
                AllDistances = two_rE*DLdDL_DS2dDS
                EinsteinAngle = (4*Goverc2*M*(DS[Ki:] - DL[:Kf])/(DS[Ki:]*DL[:Kf]))**0.5
                # 3d
                # meanVtApp:  Mean angular velocity difference
                # sigma2vtApp:  Variance in angular velocity difference
                # RMSvtapp:  RMS angular velocity difference (mean and variance combined)
                # VolumeElements:  Combine angular-only and distance-only terms
                # GammaK:  Combine all.  GammaK gives the lensing contribution at each location for this staggering of this lens and source population
                meanVtAppl = vrelappl_S[Ki:,:,:] - vrelappl_L[:Kf,:,:]
                sigma2vtAppl = sigmaappl2_S[Ki:,:,:] + sigmaappl2_L[:Kf,:,:]
                sigma2vtAppb = sigmaappb2_S[Ki:,:,:] + sigmaappb2_L[:Kf,:,:]
                RMSvtapp = (meanVtAppl**2 + sigma2vtAppl + sigma2vtAppb)**0.5
                VolumeElements = numpy.multiply(AllDistances[:,numpy.newaxis,numpy.newaxis], dOmega[numpy.newaxis,:,:])
                GammaK = nS[Ki:,:,:]*nL[:Kf,:,:]*RMSvtapp*VolumeElements
                # Sum over l, then d, to have b-dependence
                GammaKldsummed = (GammaK.sum(axis=1)).sum(axis=0)
                bHist = bHist + GammaKldsummed
                # Sum over b, then l, to have D-dependence
                GammaKblsummed = (GammaK.sum(axis=2)).sum(axis=1)
                DSHist[Ki:] = DSHist[Ki:] + GammaKblsummed
                DLHist[:Kf] = DLHist[:Kf] + GammaKblsummed
                # Final summation to get total rate from this staggering/source/lens integrand term
                ijkcontrib = GammaKblsummed.sum()
                ijcontrib = ijcontrib + ijkcontrib
                if timescales == True:
                    for MCindex in range(0, NtMult):
                        # Generate a draw from the velocity distribution.  Units are radians per year.
                        randomVl = numpy.random.normal(loc=meanVtAppl, scale=sigma2vtAppl**0.5)
                        randomVb = numpy.random.normal(loc=0, scale=sigma2vtAppb**0.5)
                        randomV = (randomVl**2 + randomVb**2)**0.5
                        # Convert that to a timescale.
                        randomT = numpy.divide(EinsteinAngle[:, numpy.newaxis, numpy.newaxis], randomV)
                        # Record the contribution to the lensing rate GammaK, which will act as a probability weighting later, and the randomT generated.  Convert both to a flattened-out form because their physical placement doesn't matter after this.
                        t1 = timescaleMasterIndex
                        t2 = t1 + randomT.size
                        timescaleArray[t1:t2] = randomT.ravel()
                        timescaleWeight[t1:t2] = GammaK.ravel()
#                        print "i, j, k, MCindex, timescaleMasterIndex, randomT.size, t2, NtGiven"
#                        print i, j, k, MCindex, timescaleMasterIndex, randomT.size, t2, NtGiven
                        timescaleMasterIndex = t2
            MassHist[j] = MassHist[j] + ijcontrib
            icontrib = icontrib + ijcontrib
        lensrate = lensrate + icontrib
        MasterbHist.append(numpy.array(bHist)) # Unlike others, it is important to collect this quantity for each magnitude
        MvHist[i] = icontrib
#        print "For a width of", dMv, "around a source magnitude of", Mv, "I found a lensing rate contribution of", icontrib
        print asctime()
        # One last pass to catch anything that made it all the way through
        if i == (len(appearance) - 1):
            lensrates[SourceMagIndex] = lensrate
            MasterMvHist.append(numpy.array(MvHist))
            MasterMassHist.append(numpy.array(MassHist))
            MasterDSHist.append(numpy.array(DSHist))
            MasterDLHist.append(numpy.array(DLHist))
#            MasterbHist.append(numpy.array(bHist))
    if timescales == True:
        hichance = float(timescaleWeight.max())
        timescaleWeight = timescaleWeight/hichance
        timeChance = numpy.random.rand(NtGiven)
        chosenTimes = timescaleArray[ numpy.where( timeChance < timescaleWeight ) ]
        chosenTimes = chosenTimes*365.25 # years -> days
        print "len(timescaleArray)", len(timescaleArray)
        print "len(chosenTimes)", len(chosenTimes)
        print "chosenTimes"
        print chosenTimes
    elif timescales == False:
        chosenTimes = numpy.zeros(1)
    if len(appearance) != 1:
        for i in range(0, len(SourceMag)):
            print "For a limiting unlensed source magnitude of", str(SourceMag[i])+ ", a lensing rate of", str(lensrates[i]), "was found."
            pass
    if len(appearance) == 1:
        print "Process", codeindex, "finishing."
        print asctime()
    # Divide each histogram by the width of the bin (i.e., normalize).
    for i in range(0, len(MasterMvHist)):
        MasterMvHist[i] = MasterMvHist[i]/delMvVals[:] # Events per magnitude
        MasterMassHist[i] = MasterMassHist[i]/delMvMass[:,1] # Events per Msun
        MasterDSHist[i] = MasterDSHist[i]/dDS # Events per parsec
        MasterDLHist[i] = MasterDLHist[i]/dDL # Events per parsec
    for i in range(0, len(MasterbHist)):
        MasterbHist[i] = MasterbHist[i]/deltaba # Events per radian
    print "GM_SimpleIntegrate ran just fine, thx for asking."
    return (lensrates, MasterMvHist, MasterMassHist, MasterDSHist, MasterDLHist, MasterbHist, chosenTimes)
                


############################
# COMMANDLINE OPERATION courtesy of codemaster Rachel Street!

if __name__ == '__main__':
    # This provision is for multithreading only, really
    # The comstr will be of the form 'python thismodule parfile outfile'

    # Import required functions:
    from sys import argv
    import cPickle as cP
    
    # Read commandline argument indicating the location of the input arguments
    parfilename = argv[-2]

    # Open and unpack
    parfile = open(parfilename, 'r')
    pars = cP.load(parfile)
    parfile.close()
    for i in pars:
        cmd = "%s = pars['%s']" % (i,i)
        exec cmd

    # Call function
    ratestuple = simple(simpars, phypars, appearance, masstuple, vtuple)

    # Read commandline argument indicating the output location
    outfilename = argv[-1]

    # Pack and write
    outpars = {'simpars' : simpars,
               'phypars' : phypars,
               'ratestuple' : ratestuple}
    outfile = open(outfilename, 'w')
    cP.dump(outpars, outfile)
    outfile.close()
